package io.kwdb.sdk.internal;

import io.kwdb.sdk.Config;
import java.io.*;
import java.net.http.HttpRequest;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.security.KeyStore;
import java.security.cert.CertificateFactory;
import java.security.cert.X509Certificate;
import java.util.Base64;
import java.util.logging.Logger;
import java.util.zip.GZIPOutputStream;
import javax.net.ssl.SSLContext;
import javax.net.ssl.TrustManagerFactory;

/** Security utilities for TLS, authentication, and compression. */
public final class SecurityUtil {

  private static final Logger logger = Logger.getLogger(SecurityUtil.class.getName());

  private SecurityUtil() {
    // Utility class
  }

  /**
   * Configure HTTP request builder with security headers based on config.
   *
   * @param builder The HTTP request builder
   * @param config The configuration
   * @param hasBody Whether the request has a body (affects gzip)
   */
  public static void configureHttpSecurity(
      HttpRequest.Builder builder, Config config, boolean hasBody) {
    // Add authentication headers
    if (config.getAuthToken() != null && !config.getAuthToken().isEmpty()) {
      builder.header("Authorization", "Bearer " + config.getAuthToken());
      logger.fine("Added Bearer token authentication");
    } else if (config.getUsername() != null && config.getPassword() != null) {
      String auth = config.getUsername() + ":" + config.getPassword();
      String encodedAuth =
          Base64.getEncoder().encodeToString(auth.getBytes(StandardCharsets.UTF_8));
      builder.header("Authorization", "Basic " + encodedAuth);
      logger.fine("Added Basic authentication for user: " + config.getUsername());
    }

    // Note: gzip compression is handled separately in the body publisher
    // Content-Encoding header should be added when body is actually compressed
  }

  /**
   * Compress data using gzip.
   *
   * @param data The data to compress
   * @return Compressed data
   * @throws IOException if compression fails
   */
  public static byte[] gzip(byte[] data) throws IOException {
    ByteArrayOutputStream baos = new ByteArrayOutputStream();
    try (GZIPOutputStream gzipOut = new GZIPOutputStream(baos)) {
      gzipOut.write(data);
    }
    return baos.toByteArray();
  }

  /**
   * Compress string data using gzip.
   *
   * @param data The string to compress
   * @return Compressed data
   * @throws IOException if compression fails
   */
  public static byte[] gzip(String data) throws IOException {
    return gzip(data.getBytes(StandardCharsets.UTF_8));
  }

  /**
   * Create an SSL context with custom CA certificate if specified.
   *
   * @param caCertPath Path to CA certificate file (PEM format)
   * @return SSLContext configured with the CA certificate
   * @throws Exception if certificate loading fails
   */
  public static SSLContext createSSLContext(String caCertPath) throws Exception {
    if (caCertPath == null || caCertPath.isEmpty()) {
      // Use system default trust store
      return SSLContext.getDefault();
    }

    // Load custom CA certificate
    CertificateFactory cf = CertificateFactory.getInstance("X.509");
    X509Certificate caCert;

    try (InputStream certStream = Files.newInputStream(Paths.get(caCertPath))) {
      caCert = (X509Certificate) cf.generateCertificate(certStream);
    }

    // Create a trust store containing the CA certificate
    KeyStore trustStore = KeyStore.getInstance(KeyStore.getDefaultType());
    trustStore.load(null, null);
    trustStore.setCertificateEntry("ca-cert", caCert);

    // Create trust manager factory with our trust store
    TrustManagerFactory tmf =
        TrustManagerFactory.getInstance(TrustManagerFactory.getDefaultAlgorithm());
    tmf.init(trustStore);

    // Create SSL context
    SSLContext sslContext = SSLContext.getInstance("TLS");
    sslContext.init(null, tmf.getTrustManagers(), null);

    logger.info("Created SSL context with custom CA certificate from: " + caCertPath);
    return sslContext;
  }

  /**
   * Get the scheme (http or https) based on TLS configuration.
   *
   * @param config The configuration
   * @return "https" if TLS is enabled, "http" otherwise
   */
  public static String getScheme(Config config) {
    return config.isTls() ? "https" : "http";
  }

  /**
   * Build URL for ILP endpoint with proper scheme.
   *
   * @param config The configuration
   * @return Complete URL for the write endpoint
   */
  public static String buildIlpUrl(Config config) {
    String scheme = getScheme(config);
    int port = config.getHttpPort() > 0 ? config.getHttpPort() : 9091;
    return String.format("%s://%s:%d/write", scheme, config.getHost(), port);
  }

  /**
   * Apply security configuration to HTTP client builder.
   *
   * @param builder The HTTP client builder
   * @param config The configuration
   * @return The configured builder
   */
  public static java.net.http.HttpClient.Builder configureHttpClient(
      java.net.http.HttpClient.Builder builder, Config config) {

    if (config.isTls()) {
      try {
        SSLContext sslContext = createSSLContext(config.getCaCertPath());
        builder.sslContext(sslContext);
      } catch (Exception e) {
        logger.warning("Failed to configure custom SSL context: " + e.getMessage());
        // Fall back to default SSL context
      }
    }

    return builder;
  }
}
